package com.gitee.zhziyun.cnfsystem.utils.FileUtil;

import com.gitee.zhziyun.cnfsystem.utils.XXHashUtil;

import java.io.*;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

/**
 * <p>多线程文件接收工具（支持断点续传、进度控制）</p>
 * <p>功能特性：</p>
 * <p>1. 自动识别单线程/多线程传输模式</p>
 * <p>2. 实时进度回调</p>
 * <p>3. 支持暂停/恢复操作</p>
 * <p>4. 自动文件校验和重组</p>
 */
public class ThreadedFileReceiver {
    // 常量定义
    private static final int BLOCK_SIZE = 1024 * 4; // 必须与发送端一致
    private static final int MAX_RETRIES = 3; // 最大重试次数

    // 传输控制状态
    private volatile boolean paused = false;
    private volatile boolean cancelled = false;

    // 配置参数
    private final int port;
    private final String saveDir;
    private final ProgressListener progressListener;
    private final Map<Long, File> tempFiles = new ConcurrentHashMap<>();

    /**
     * 进度监听接口，用于接收文件传输过程中的各种事件
     */
    public interface ProgressListener {
        default void onProgress(long current, long total) {}
        default void onPaused() {}
        default void onResumed() {}
        default void onCompleted() {}
        default void onError(Exception e) {}
    }

    /**
     * 构造方法，创建一个新的文件接收器实例
     * 
     * @param port 监听端口，等待发送端连接
     * @param saveDir 文件保存目录，用于存储接收到的文件
     * @param listener 进度监听器，用于监听传输过程
     */
    public ThreadedFileReceiver(int port, String saveDir, ProgressListener listener) {
        this.port = port;
        this.saveDir = saveDir;
        this.progressListener = listener;
    }

    // 主要控制方法
    public void start() { new Thread(this::doReceive).start(); }
    public void pause() { paused = true; progressListener.onPaused(); }
    public void resume() { paused = false; progressListener.onResumed(); }
    public void cancel() { cancelled = true; }

    /**
     * 核心接收方法，处理整个文件接收流程
     */
    private void doReceive() {
        try (ServerSocket serverSocket = new ServerSocket(port)) {
            Socket socket = serverSocket.accept();
            DataInputStream dis = new DataInputStream(socket.getInputStream());

            // 接收文件元信息
            String fileHash = dis.readUTF();
            String fileName = dis.readUTF();
            long fileSize = dis.readLong();
            boolean isMultiThread = dis.readBoolean();

            File outputDir = new File(saveDir, fileHash);
            // 如果目录已存在，递归删除整个目录
            if (outputDir.exists()) {
                deleteDirectoryRecursively(outputDir);
            }

            if (!outputDir.mkdirs()) {
                throw new IOException("无法创建目录: " + outputDir.getAbsolutePath());
            }

            if (isMultiThread) {
                multiThreadReceive(dis, outputDir, fileSize);
            } else {
                singleThreadReceive(dis, outputDir, fileSize);
            }

            // 合并文件
            File outputFile = new File(saveDir, fileName);
            FileConcatenator.concatenateTempFiles(outputDir.getAbsolutePath(),
                    outputFile.getAbsolutePath());

            // 校验文件完整性
            if (!validateFile(outputFile, fileHash)) {
                throw new IOException("文件校验失败");
            }

            // 删除整个临时目录（包含所有.tmp文件）
            deleteDirectoryRecursively(outputDir);

            progressListener.onCompleted();
        } catch (Exception e) {
            progressListener.onError(e);
        }
    }

    /**
     * 单线程接收实现，顺序接收整个文件
     * 
     * @param dis 数据输入流，来自发送端的连接
     * @param outputDir 输出目录，存储接收的临时文件
     * @param fileSize 文件大小，用于进度控制
     * @throws IOException 如果接收过程中发生IO异常
     */
    private void singleThreadReceive(DataInputStream dis, File outputDir,
                                     long fileSize) throws IOException {
        try (FileOutputStream fos = new FileOutputStream(new File(outputDir, "0.tmp"))) {
            byte[] buffer = new byte[BLOCK_SIZE];
            long received = 0;
            int retries = 0;

            while (!cancelled && received < fileSize && retries < MAX_RETRIES) {
                if (paused) {
                    Thread.sleep(100);
                    continue;
                }

                try {
                    int bytesToRead = (int) Math.min(BLOCK_SIZE, fileSize - received);
                    int bytesRead = dis.read(buffer, 0, bytesToRead);
                    if (bytesRead == -1) break;

                    fos.write(buffer, 0, bytesRead);
                    received += bytesRead;
                    progressListener.onProgress(received, fileSize);
                    retries = 0; // 重置重试计数
                } catch (IOException e) {
                    retries++;
                    Thread.sleep(1000);
                }
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    /**
     * 多线程接收实现，同时接收文件的不同部分并最终合并
     * 
     * @param dis 数据输入流，来自发送端的连接
     * @param outputDir 输出目录，存储接收的分块文件
     * @param fileSize 文件总大小，用于进度控制
     * @throws Exception 如果接收过程中发生任何异常
     */
    private void multiThreadReceive(DataInputStream dis, File outputDir,
                                    long fileSize) throws Exception {
        ExecutorService executor = Executors.newFixedThreadPool(4);
        long chunkCount = (fileSize + BLOCK_SIZE - 1) / BLOCK_SIZE;
        AtomicLong receivedChunks = new AtomicLong(0);
        int retries = 0;
        Set<Long> failedChunks = Collections.synchronizedSet(new HashSet<>());

        try {
            while (!cancelled && receivedChunks.get() < chunkCount &&
                    retries < MAX_RETRIES) {
                if (paused) {
                    Thread.sleep(100);
                    continue;
                }

                try {
                    long chunkIndex = dis.readLong();
                    int chunkSize = dis.readInt();
                    
                    // 读取8个字节的哈希值
                    byte[] expectedHashBytes = new byte[8];
                    dis.readFully(expectedHashBytes);
                    
                    byte[] buffer = new byte[chunkSize];
                    dis.readFully(buffer);

                    // 验证分块哈希
                    byte[] actualHashBytes = XXHashUtil.calculateHash(buffer);
                    if (!Arrays.equals(expectedHashBytes, actualHashBytes)) {
                        failedChunks.add(chunkIndex);
                        retries++;
                        continue;
                    }

                    // 使用final变量避免lambda捕获问题
                    final byte[] chunkData = buffer.clone();
                    final long finalChunkIndex = chunkIndex;
                    executor.submit(() -> {
                        try {
                            File chunkFile = new File(outputDir, chunkIndex + ".tmp");
                            try (FileOutputStream fos = new FileOutputStream(chunkFile)) {
                                fos.write(chunkData);
                            }
                            // 使用putIfAbsent确保线程安全
                            tempFiles.putIfAbsent(finalChunkIndex, chunkFile);
                            long currentChunk = receivedChunks.incrementAndGet();
                            // 更精确的进度计算
                            progressListener.onProgress(
                                    currentChunk * (long)BLOCK_SIZE, fileSize);
                            
                            // 如果之前失败过，现在成功了，从失败列表中移除
                            failedChunks.remove(finalChunkIndex);
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    });
                    retries = 0;
                } catch (IOException e) {
                    retries++;
                    Thread.sleep(1000);
                }
            }
            
            // 重新传输失败的分块
            if (!failedChunks.isEmpty()) {
                progressListener.onProgress(receivedChunks.get() * BLOCK_SIZE, fileSize);
                // 递归调用自身来重传失败的分块
                multiThreadReceive(dis, outputDir, fileSize);
            }
        } finally {
            executor.shutdown();
            executor.awaitTermination(1, TimeUnit.HOURS);
        }
    }

    // 辅助方法
    /**
     * 验证文件完整性，通过比较哈希值确认文件未损坏
     * 
     * @param file 要验证的文件
     * @param expectedHash 预期的哈希值
     * @return 如果文件哈希匹配返回true，否则返回false
     * @throws IOException 如果读取文件时发生IO异常
     */
    private boolean validateFile(File file, String expectedHash) throws IOException {
        return XXHashUtil.calculateFileHashToString(file).equals(expectedHash);
    }

    /**
     * 递归删除目录及其所有内容（包括自身）
     * 
     * @param dir 要删除的目录
     * @throws IOException 如果删除过程中发生IO异常
     */
    private static void deleteDirectoryRecursively(File dir) throws IOException {
        if (dir.isDirectory()) {
            File[] files = dir.listFiles();
            if (files != null) {
                for (File file : files) {
                    deleteDirectoryRecursively(file);
                }
            }
        }
        if (!dir.delete()) {
            throw new IOException("无法删除文件或目录: " + dir.getAbsolutePath());
        }
    }
}